import json

from openai import OpenAI
from tenacity import retry, wait_random_exponential, stop_after_attempt
from termcolor import colored

GPT_MODEL = "gpt-4o"
client = OpenAI()
client.api_key = "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"

def get_n_day_weather_forecast(city, tep_format, days):
    # return f"{city}未来{days}天的天气 {tep_format} 非常棒的"
    # return "very bad and There may be heavy rain"
    # print(f'this is user  input {user_input}\n')
    return "very bad and There may be heavy rain"

@retry(wait=wait_random_exponential(multiplier=1, max=40), stop=stop_after_attempt(3))
def chat_completion_request(messages, tools=None, tool_choice=None, model=GPT_MODEL):
    try:
        response = client.chat.completions.create(
            model=model,
            messages=messages,
            tools=tools,
            tool_choice=tool_choice,
        )
        return response
    except Exception as e:
        print("Unable to generate ChatCompletion response")
        print(f"Exception: {e}")
        return e


def pretty_print_conversation(messages):
    role_to_color = {
        "system": "red",
        "user": "green",
        "assistant": "blue",
        "function": "magenta",
    }

    for message in messages:
        if message["role"] == "system":
            print(colored(f"system: {message['content']}\n", role_to_color[message["role"]]))
        elif message["role"] == "user":
            print(colored(f"user: {message['content']}\n", role_to_color[message["role"]]))
        elif message["role"] == "assistant" and message.get("function_call"):
            print(colored(f"assistant: {message['function_call']}\n", role_to_color[message["role"]]))
        elif message["role"] == "assistant" and not message.get("function_call"):
            print(colored(f"assistant: {message['content']}\n", role_to_color[message["role"]]))
        elif message["role"] == "function":
            print(colored(f"function ({message['name']}): {message['content']}\n", role_to_color[message["role"]]))


if __name__ == '__main__':

    tools = [
        {
            "type": "function",
            "function": {
                "name": "get_current_weather",
                "description": "Get the current weather",
                "parameters": {
                    "type": "object",
                    "properties": {

                        "location": {
                            "type": "string",
                            "description": "The city and state, e.g. San Francisco, CA",
                        },

                        "format": {
                            "type": "string",
                            "enum": ["celsius", "fahrenheit"],
                            "description": "The temperature unit to use. Infer this from the users location.",
                        },
                    },
                    "required": ["location", "format"],
                },
            }
        },
        {
            "type": "function",
            "function": {
                "name": "get_n_day_weather_forecast",
                "description": "Get an N-day weather forecast",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "The city and state, e.g. San Francisco, CA",
                        },
                        "format": {
                            "type": "string",
                            "enum": ["celsius", "fahrenheit"],
                            "description": "The temperature unit to use. Infer this from the users location.",
                        },
                        "num_days": {
                            "type": "integer",
                            "description": "The number of days to forecast",
                        }
                    },
                    "required": ["location", "format", "num_days"]
                },
            }
        },
    ]

    # 1.当没有提供地址，会提示要求提供地址
    # messages = []
    # messages.append({"role": "system",
    #                  "content": "Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous."})
    # messages.append({"role": "user", "content": "What's the weather like today"})
    # chat_response = chat_completion_request(
    #     messages, tools=tools
    # )
    # assistant_message = chat_response.choices[0].message
    # messages.append(assistant_message)
    # print(assistant_message)

    # 2.当有提供地址，会提示要求提供地址
    # messages = []
    # messages.append({"role": "system",
    #                  "content": "Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous."})
    # messages.append({"role": "user", "content": "What's the weather like today"})
    # messages.append({"role": "user", "content": "I'm in Glasgow, Scotland."})
    # chat_response = chat_completion_request(
    #     messages, tools=tools
    # )
    # assistant_message = chat_response.choices[0].message
    # messages.append(assistant_message)
    # print(assistant_message)

    # 3.天气预报预测
    # messages = []
    # messages.append({"role": "system",
    #                  "content": "Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous."})
    # messages.append(
    #     {"role": "user", "content": "what is the weather going to be like in Glasgow, Scotland over the next x days"})
    # messages.append({"role": "user", "content": "5 days"})
    # chat_response = chat_completion_request(
    #     messages, tools=tools
    # )
    # assistant_message = chat_response.choices[0].message
    # messages.append(assistant_message)
    # print(assistant_message)

    # 4.强制使用自定义函数
    messages = []

    messages.append({"role": "system","content": "Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous."})
    messages.append({"role": "user", "content": "what is the weather going to be like in Glasgow, Scotland over the next 5 days"})
    # messages.append({"role": "user", "content": "5 days"})

    chat_response = chat_completion_request(
        messages, tools=tools, tool_choice={"type": "function", "function": {"name": "get_n_day_weather_forecast"}}
    )

    response_message = chat_response.choices[0].message
    print(chat_response.choices[0].message)
    messages.append(response_message) #//很重要必须，加上这个message，否则Ai无法得到最终的结果

    tool_calls = response_message.tool_calls

    if tool_calls:
        tool_call_id = tool_calls[0].id
        tool_function_name = tool_calls[0].function.name
        paramsMap = tool_calls[0].function.arguments
        my_dict = json.loads(paramsMap)
        print("-------------------------------")
        for key, value in my_dict.items():
            print("参数-->Key:", key, ", 参数 Value:", value)
        print("-------------------------------")
        print(f"调用了自定义函数--->>{tool_function_name}")
        results = get_n_day_weather_forecast(my_dict['location'], my_dict['format'],my_dict['num_days'])
        messages.append({
            "role": "tool",
            "tool_call_id": tool_call_id,
            "name": tool_function_name,
            "content": results
        })

        # Step 4: Invoke the chat completions API with the function response appended to the messages list
        # Note that messages with role 'tool' must be a response to a preceding message with 'tool_calls'
        model_response_with_function_call = client.chat.completions.create(
            model=GPT_MODEL,
            messages=messages,
        )  # get a new response from the model where it can see the function response
        print("最终结果---->> ",model_response_with_function_call.choices[0].message.content)

    else:
        print("没有调用自定函数")



    # 5.不强制使用，结果可能不会调用自定义函数
    # if we don't force the model to use get_n_day_weather_forecast it may not
    # messages = []
    # messages.append({"role": "system",
    #                  "content": "Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous."})
    # messages.append({"role": "user", "content": "Give me a weather report for Toronto, Canada."})
    # chat_response = chat_completion_request(
    #     messages, tools=tools
    # )
    # chat_response.choices[0].message

    # 6.强制不使用自定义函数
    # messages = []
    # messages.append({"role": "system",
    #                  "content": "Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous."})
    # messages.append({"role": "user", "content": "Give me the current weather (use Celcius) for Toronto, Canada."})
    # chat_response = chat_completion_request(
    #     messages, tools=tools, tool_choice="none"
    # )
    # chat_response.choices[0].message

    # 7.并行函数调用
    # messages = []
    # messages.append({"role": "system",
    #                  "content": "Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous."})
    # messages.append({"role": "user",
    #                  "content": "what is the weather going to be like in San Francisco and Glasgow over the next 4 days"})
    # chat_response = chat_completion_request(
    #     messages, tools=tools, model=GPT_MODEL
    # )
    # assistant_message = chat_response.choices[0].message.tool_calls
    # print(assistant_message)
